#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @FileName  : beautify.py
# @Time      : 2025/5/26
# @Author    : 雨霓同学
# @Function  : YOLO 检测结果美化绘制（圆角标签、圆角检测框、中英文支持优化，含特殊圆角处理）

from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np
from collections import OrderedDict
import os

# ======================= 全局配置 ==========================
FONT_PATH = r"D:\project\Vehicle_inspection\roadmaincode\LXGWWenKai-Bold.ttf"
FONT_SIZE = 20  # 默认字体大小（基于720p）
LINE_WIDTH = 4  # 默认线宽
LABEL_PADDING = (10, 10)  # 默认标签内边距（水平，垂直）
RADIUS = 8  # 默认圆角半径
TEXT_COLOR = (0, 0, 0)  # 默认文本颜色（BGR，黑色）

LABEL_MAPPING = {
    'pedestrian': "行人",
    'people': "人",
    'bicycle': "自行车",
    'car': "汽车",
    'van': "轻型商用车",
    'truck': "卡车",
    'tricycle': "三轮车",
    'awning-tricycle': "遮阳三轮车",
    'bus': "大巴",
    'motor': "摩托车",
    'others': "其他"
}
COLOR_MAPPING = {
    "safety_helmet": (50, 205, 50),  # 绿色 (BGR)
    "reflective_vest": (255, 255, 0),  # 黄色 (BGR)
    "person": (65, 105, 255),  # 橙色 (BGR)
    "head": (0, 0, 255),  # 红色 (BGR)
    "ordinary_clothes": (226, 43, 138)  # 紫色 (BGR)
}
text_size_cache = OrderedDict()


# ======================= 美化辅助函数（基于OpenCV） ======================
def draw_filled_rounded_rect(image_np, pt1, pt2, color_bgr, radius,
                            top_left_round=True, top_right_round=True,
                            bottom_left_round=True, bottom_right_round=True):
    """使用 OpenCV 绘制颜色填充的圆角矩形，可控制每个角的圆角状态"""
    x1, y1 = pt1
    x2, y2 = pt2
    thickness = -1
    cv2.rectangle(image_np, (x1 + (radius if top_left_round else 0), y1),
                (x2 - (radius if top_right_round else 0), y2), color_bgr, thickness)
    cv2.rectangle(image_np, (x1, y1 + (radius if top_left_round else 0)),
                (x2, y2 - (radius if bottom_left_round else 0)), color_bgr, thickness)
    if top_left_round:
        cv2.circle(image_np, (x1 + radius, y1 + radius), radius, color_bgr, thickness, cv2.LINE_AA)
    if top_right_round:
        cv2.circle(image_np, (x2 - radius, y1 + radius), radius, color_bgr, thickness, cv2.LINE_AA)
    if bottom_left_round:
        cv2.circle(image_np, (x1 + radius, y2 - radius), radius, color_bgr, thickness, cv2.LINE_AA)
    if bottom_right_round:
        cv2.circle(image_np, (x2 - radius, y2 - radius), radius, color_bgr, thickness, cv2.LINE_AA)


def draw_bordered_rounded_rect(image_np, pt1, pt2, color_bgr, thickness, radius,
                            top_left_round=True, top_right_round=True,
                            bottom_left_round=True, bottom_right_round=True):
    """使用 OpenCV 绘制带边框的圆角矩形，可控制每个角的圆角状态"""
    x1, y1 = pt1
    x2, y2 = pt2
    line_type = cv2.LINE_AA
    if top_left_round or top_right_round:
        cv2.line(image_np, (x1 + (radius if top_left_round else 0), y1),
                (x2 - (radius if top_right_round else 0), y1), color_bgr, thickness, line_type)
    if bottom_left_round or bottom_right_round:
        cv2.line(image_np, (x1 + (radius if bottom_left_round else 0), y2),
                (x2 - (radius if bottom_right_round else 0), y2), color_bgr, thickness, line_type)
    if top_left_round or bottom_left_round:
        cv2.line(image_np, (x1, y1 + (radius if top_left_round else 0)),
                (x1, y2 - (radius if bottom_left_round else 0)), color_bgr, thickness, line_type)
    if top_right_round or bottom_right_round:
        cv2.line(image_np, (x2, y1 + (radius if top_right_round else 0)),
                (x2, y2 - (radius if bottom_right_round else 0)), color_bgr, thickness, line_type)
    if top_left_round:
        cv2.ellipse(image_np, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color_bgr, thickness, line_type)
    if top_right_round:
        cv2.ellipse(image_np, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color_bgr, thickness, line_type)
    if bottom_left_round:
        cv2.ellipse(image_np, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color_bgr, thickness, line_type)
    if bottom_right_round:
        cv2.ellipse(image_np, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color_bgr, thickness, line_type)

# ======================= 文本和缓存辅助函数（基于Pillow） ======================
def draw_text_pil(image, text, position, font_obj, fill_bgr=TEXT_COLOR):
    """使用 PIL 绘制文本到 OpenCV 图像上"""
    image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)
    fill_rgb = (fill_bgr[2], fill_bgr[1], fill_bgr[0])
    draw.text(position, text, font=font_obj, fill=fill_rgb)
    return cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)

def preload_cache(font_sizes=[20, 22, 30, 44]):
    """预缓存中英文标签尺寸"""
    global text_size_cache
    text_size_cache.clear()
    for size in font_sizes:
        try:
            font = ImageFont.truetype(FONT_PATH, size)
        except IOError:
            print(f"警告：无法加载字体文件 '{FONT_PATH}'，跳过字体大小 {size} 的预缓存")
            continue
        for label in LABEL_MAPPING.values():
            text = f"{label} 80.0%"
            cache_key = f"{text}_{size}"
            temp_image = Image.new('RGB', (1, 1))
            draw = ImageDraw.Draw(temp_image)
            bbox = draw.textbbox((0, 0), text, font=font)
            text_size_cache[cache_key] = (bbox[2] - bbox[0], bbox[3] - bbox[1])
        for label in LABEL_MAPPING.keys():
            text = f"{label} 80.0%"
            cache_key = f"{text}_{size}"
            temp_image = Image.new('RGB', (1, 1))
            draw = ImageDraw.Draw(temp_image)
            bbox = draw.textbbox((0, 0), text, font=font)
            text_size_cache[cache_key] = (bbox[2] - bbox[0], bbox[3] - bbox[1])


def get_text_size(text, font_obj, max_cache_size=50):
    """计算文本尺寸（带缓存，规范化置信度）"""
    label = text.split(" ")[0]
    cache_key = f"{label} 80.0%_{font_obj.size}"
    if cache_key in text_size_cache:
        return text_size_cache[cache_key]
    temp_image = Image.new('RGB', (1, 1))
    draw = ImageDraw.Draw(temp_image)
    bbox = draw.textbbox((0, 0), text, font=font_obj)
    width = bbox[2] - bbox[0]
    height = bbox[3] - bbox[1]
    text_size_cache[cache_key] = (width, height)
    if len(text_size_cache) > max_cache_size:
        text_size_cache.popitem(last=False)
    return (width, height)


# ======================= 核心绘制函数 ======================
def custom_plot(
    image,
    boxes,
    confs,
    labels,
    use_chinese_mapping=True,
    beautify=True,
    font_size=FONT_SIZE,
    line_width=LINE_WIDTH,
    label_padding=LABEL_PADDING,
    radius=RADIUS,
    text_color=TEXT_COLOR
):
    """绘制检测框和标签"""
    result_image = image.copy()
    ref_dim = 720
    current_short_dim = min(image.shape[0], image.shape[1])
    scale_factor = current_short_dim / ref_dim if ref_dim > 0 else 1.0
    scaled_font_size = max(10, int(font_size * scale_factor))
    scaled_line_width = max(2, int(line_width * scale_factor))
    scaled_padding_x = int(label_padding[0] * scale_factor)
    scaled_padding_y = int(label_padding[1] * scale_factor)
    scaled_radius = max(1, int(radius * scale_factor))

    try:
        font_pil = ImageFont.truetype(FONT_PATH, scaled_font_size)
    except OSError:
        print(f"错误：无法加载字体文件 '{FONT_PATH}'。将使用Pillow默认字体。")
        font_pil = ImageFont.load_default()

    for box, conf, label_key in zip(boxes, confs, labels):
        x1, y1, x2, y2 = map(int, box)
        color_bgr = COLOR_MAPPING.get(label_key, (0, 255, 0))
        # 标签语言
        if beautify and use_chinese_mapping:
            display_label = LABEL_MAPPING.get(label_key, label_key)
        else:
            display_label = label_key  # 非美化强制英文
        label_text_full = f"{display_label} {conf * 100:.1f}%"

        # 计算标签框尺寸和位置
        text_width, text_height = get_text_size(label_text_full, font_pil)
        label_box_actual_width = text_width + 2 * scaled_padding_x
        if label_box_actual_width < 2 * scaled_radius:
            label_box_actual_width = 2 * scaled_radius
        label_box_actual_height = text_height + 2 * scaled_padding_y

        # 左对齐
        label_box_x_min = int(x1 - scaled_line_width / 2.0)
        label_box_y_min = y1 - label_box_actual_height
        if label_box_y_min < 0:
            label_box_y_min = 0
        label_box_x_max = label_box_x_min + label_box_actual_width
        label_box_y_max = y1

        # 边界检查
        if label_box_x_max > image.shape[1]:
            label_box_x_max = image.shape[1]
            label_box_x_min = label_box_x_max - label_box_actual_width
            if label_box_x_min < 0:
                label_box_x_min = 0
                label_box_x_max = label_box_x_min + label_box_actual_width

        # 特殊圆角规则
        # 判断标签框是否比检测框宽
        is_label_wider_than_box = label_box_actual_width > (x2 - x1)
        label_top_left_round = True
        label_top_right_round = True
        label_bottom_left_round = False
        label_bottom_right_round = (is_label_wider_than_box and beautify)
        det_top_left_round = False
        det_top_right_round = not is_label_wider_than_box
        det_bottom_left_round = True
        det_bottom_right_round = True

        if beautify:
            # 绘制检测框
            draw_bordered_rounded_rect(result_image, (x1, y1), (x2, y2),
                                    color_bgr, scaled_line_width, scaled_radius,
                                    det_top_left_round, det_top_right_round,
                                    det_bottom_left_round, det_bottom_right_round)
            # 绘制标签框
            draw_filled_rounded_rect(result_image, (label_box_x_min, label_box_y_min),
                                    (label_box_x_max, label_box_y_max),
                                    color_bgr, scaled_radius,
                                    label_top_left_round, label_top_right_round,
                                    label_bottom_left_round, label_bottom_right_round)
            # 绘制文本
            text_x = label_box_x_min + (label_box_actual_width - text_width) // 2
            text_y = label_box_y_min + (label_box_actual_height - text_height) // 2
            result_image = draw_text_pil(result_image, label_text_full, (text_x, text_y), font_pil, text_color)
    return result_image
